package edu.berkeley.nlp.assignments;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import edu.berkeley.nlp.langmodel.LanguageModel;
import edu.berkeley.nlp.util.Counter;
import edu.berkeley.nlp.util.CounterMap;

public class KatzBigramLanguageModel implements LanguageModel {

	static final String STOP = "</S>";

	
	double total = 0.0;
	// for unigram, c(w')
	Counter<String> wordCounter = new Counter<String>();

	double totalBigram = 0.0;
	// for bigram, c*(w,w')
    CounterMap<String, String> bigramCounterMap = new CounterMap<String, String>();

//    CounterMap<String, String> bigramKatzProb = new CounterMap<String, String>();
    
    Counter<String> katzAlpha = new Counter<String>();
    
    // katz constant
    int k = 3; 
    
    
	public List<String> generateSentence() {
		// TODO Auto-generated method stub
		return null;
	}

	public double getSentenceProbability(List<String> sentence) {
		// TODO Auto-generated method stub
		return 0;
	}

	public void KatzDiscount() {
		for (String wP: bigramCounterMap.keySet()) {
			
			Counter<String> counterFromAKey = bigramCounterMap.getCounter(wP);
			
			// FROM THIS ON, count is in double type but integer values
			// get largest count
			Double largestCount = 0.0, count;
			for (String w: counterFromAKey.keySet()) {
				if ((count = counterFromAKey.getCount(w)) > largestCount) {
					largestCount = count;
				}
			}
			
			//System.out.println("largest count " + largestCount );

			// count n0, n1,...,nk
			int arraySize = Math.max(k + 2, largestCount.intValue() + 2);
			int n[] = new int[arraySize];
			for (int i = 0; i < n.length; i++)
				n[i] = 0;
			for (String w: counterFromAKey.keySet()) {
				count = counterFromAKey.getCount(w);
				n[count.intValue()] ++;
			}
			
			
			double subtotal = counterFromAKey.totalCount();
			double zeroCountBigramTotalInUnigram = 0.0;
			// Katz discount: r* = (r+1) * n_{r+1}/n_r when r <= k => dr => dr * r
			// also count the distrubtion of zero-count bigrams
			for (String w: wordCounter.keySet()) {
				count = counterFromAKey.getCount(w);
				
				if ( counterFromAKey.containsKey(w) && count <= k) {
					if (n[1+count.intValue()] == 0)
						System.out.println("zero");
					
					// FROM this on, count is in double values
					
					// original good turing
					//count = (count+1.0)*n[1+count.intValue()]/n[count.intValue()];
					
					// good turing with k value, !!! zero values, need to smooth it!
					double rStar = (count+1.0)*n[1+count.intValue()]/n[count.intValue()];
				/*	if (k+1 > largestCount)
						count *= ((rStar / count) - largestCount * n[largestCount.intValue()] / n[1]) / (1 - (largestCount) * n[largestCount.intValue()] / n[1]);
					else*/	
						count *= ((rStar / count) - (k + 1) * n[k+1] / n[1]) / (1 - (k + 1) * n[k+1] / n[1]);
					
					bigramCounterMap.setCount(wP, w, count);
					
					
				} else if (!counterFromAKey.containsKey(w)) {
					zeroCountBigramTotalInUnigram += wordCounter.getCount(w);
				}
			}
			
			
			double dominator = (subtotal - counterFromAKey.totalCount()) / subtotal;
			System.out.println(zeroCountBigramTotalInUnigram+ "total bigram count with word [" + wP + "] discount from " + subtotal + " to " + counterFromAKey.totalCount() + " with " + dominator);

			double denominator = zeroCountBigramTotalInUnigram/ total;
			katzAlpha.setCount(wP, dominator/denominator);
			/*
			// katz smoothing
			double alpha = dominator / denominator;
			for (String w: wordCounter.keySet()) {
				if (!counterFromAKey.containsKey(w)) {
					bigramCounterMap.setCount(wP, w, alpha * wordCounter.getCount(w)/ total);
				}
			}
			System.out.println("total bigram count with word [" + wP + "] is " + counterFromAKey.totalCount() + " at final");
			*/
		}
	}
	
	
	
	public void originalProb() {
		for (String wP: bigramCounterMap.keySet()) {

			Counter<String> counterFromAKey = bigramCounterMap.getCounter(wP);
			
			double count = 0.0;
			for (String w: counterFromAKey.keySet()) {
				count = counterFromAKey.getCount(w);
				
				double prob = count / totalBigram;
				
//				System.out.println("bigram prob from " + count / totalBigram + " to " + prob);
//				bigramKatzProb.setCount(wP, w, prob);
			}
		}
	}
	
	
	public KatzBigramLanguageModel(Collection<List<String>> sentenceCollection) {
		for (List<String> sentence : sentenceCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			stoppedSentence.add(STOP);
			String lastWord = null;
			for (String word : stoppedSentence) {
				wordCounter.incrementCount(word, 1.0);
				if (lastWord != null) {
					bigramCounterMap.incrementCount(lastWord, word, 1.0);
				}
				lastWord = word;
			}
		}
		total = wordCounter.totalCount();
		totalBigram = bigramCounterMap.totalCount();
		
		originalProb();
//		System.out.println("total bigram prob before Katz " + bigramKatzProb.totalCount());
		
		KatzDiscount();
		
		System.out.println("total discounted from " + totalBigram + " to " + bigramCounterMap.totalCount() );
		
		
	}
	
	
}
